# AE聚类分析主脚本
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.datasets import EMNIST, CIFAR10
from torch.utils.data import DataLoader
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score, adjusted_rand_score
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import silhouette_samples

# 可视化错误聚类样本
def visualize_misclustered_images(features, labels, preds, images, class_names=None, max_show=25):
    mismatches = (labels != preds)
    mismatched_idx = np.where(mismatches)[0]
    print(f"总共发现 {len(mismatched_idx)} 个预测错误样本")

    plt.figure(figsize=(10, 10))
    for i, idx in enumerate(mismatched_idx[:max_show]):
        plt.subplot(5, 5, i+1)
        plt.imshow(images[idx][0], cmap='gray')
        plt.axis('off')
        title = f"T:{labels[idx]} P:{preds[idx]}"
        plt.title(title, fontsize=8)
    plt.suptitle("Misclustered Samples (True vs Pred)", fontsize=12)
    plt.tight_layout()
    plt.show()

# 聚类标签 vs 真实标签的双重 t-SNE 可视化
def visualize_dual_tsne(features, labels, preds):
    tsne = TSNE(n_components=2, random_state=42)
    reduced = tsne.fit_transform(features)

    plt.figure(figsize=(12, 5))
    plt.subplot(1, 2, 1)
    plt.scatter(reduced[:, 0], reduced[:, 1], c=labels, cmap='tab20', s=5)
    plt.title("t-SNE by True Labels")
    plt.colorbar()

    plt.subplot(1, 2, 2)
    plt.scatter(reduced[:, 0], reduced[:, 1], c=preds, cmap='tab20', s=5)
    plt.title("t-SNE by Cluster Labels")
    plt.colorbar()

    plt.tight_layout()
    plt.show()

# 可视化轮廓系数直方图
def visualize_silhouette_distribution(features, preds):
    sil_samples = silhouette_samples(features, preds)
    plt.figure(figsize=(8, 4))
    plt.hist(sil_samples, bins=50, color='skyblue', edgecolor='k')
    plt.title("Silhouette Coefficient Distribution")
    plt.xlabel("Silhouette Coefficient")
    plt.ylabel("Frequency")
    plt.grid(True)
    plt.show()



# ========== 定义 AE 模型 ==========
class AutoEncoder(nn.Module):
    def __init__(self, latent_dim=64):
        super(AutoEncoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, latent_dim)
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid(),
            nn.Unflatten(1, (1, 28, 28))
        )

    def forward(self, x):
        z = self.encoder(x)
        out = self.decoder(z)
        return out

# ========== 数据加载函数 ==========
def load_emnist():
    transform = transforms.Compose([transforms.ToTensor()])
    train = EMNIST(root="./data", split="balanced", train=True, download=True, transform=transform)
    test = EMNIST(root="./data", split="balanced", train=False, download=True, transform=transform)
    return train, test

# ========== AE训练函数 ==========
def train_autoencoder(model, dataset, epochs=10):
    dataloader = DataLoader(dataset, batch_size=128, shuffle=True)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.MSELoss()
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for x, _ in dataloader:
            x = x.cuda()
            out = model(x)
            loss = criterion(out, x)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

# ========== 特征提取函数 ==========
def extract_features(model, dataset):
    dataloader = DataLoader(dataset, batch_size=128)
    model.eval()
    features, labels = [], []
    with torch.no_grad():
        for x, y in dataloader:
            x = x.cuda()
            z = model.encoder(x).cpu()
            features.append(z)
            labels.append(y)
    return torch.cat(features).numpy(), torch.cat(labels).numpy()

# ========== 聚类评估函数 ==========
def cluster_and_evaluate(features, labels):
    kmeans = KMeans(n_clusters=47, random_state=42)
    preds = kmeans.fit_predict(features)
    sil = silhouette_score(features, preds)
    ari = adjusted_rand_score(labels, preds)
    print(f"Silhouette Score: {sil:.4f}, ARI: {ari:.4f}")
    return preds

# ========== 可视化函数 ==========
def visualize_tsne(features, labels, title="t-SNE Visualization"):
    tsne = TSNE(n_components=2, random_state=42)
    reduced = tsne.fit_transform(features)
    plt.figure(figsize=(8,6))
    plt.scatter(reduced[:,0], reduced[:,1], c=labels, cmap='tab20', s=5)
    plt.title(title)
    plt.colorbar()
    plt.show()

# ========== 主流程 ==========
def main():
    train_set, test_set = load_emnist()

    for latent_dim in [16, 32, 64, 128]:
        print(f"\nTesting latent_dim = {latent_dim}")
        model = AutoEncoder(latent_dim=latent_dim).cuda()
        train_autoencoder(model, train_set, epochs=5)
        features, labels = extract_features(model, test_set)

        # KMeans 聚类
        preds = cluster_and_evaluate(features, labels)

        # 原有 t-SNE（pred 可视化）
        visualize_tsne(features, preds, title=f"AE latent_dim={latent_dim} (Cluster Labels)")

        # 双视角 t-SNE 可视化（聚类标签 vs 真实标签）
        visualize_dual_tsne(features, labels, preds)

        # 轮廓系数直方图
        visualize_silhouette_distribution(features, preds)

        # 聚类错误样本可视化（注意需要图片）
        test_loader = DataLoader(test_set, batch_size=128)
        all_images = []
        for x, _ in test_loader:
            all_images.append(x)
        all_images = torch.cat(all_images).numpy()
        visualize_misclustered_images(features, labels, preds, all_images)

    # 原始图像 baseline 聚类
    print("\nBaseline: 使用原始图像")
    dataloader = DataLoader(test_set, batch_size=6)
    all_x, all_y = [], []
    for x, y in dataloader:
        all_x.append(x.view(x.size(0), -1))
        all_y.append(y)
    x_raw = torch.cat(all_x).numpy()
    y_raw = torch.cat(all_y).numpy()
    preds = cluster_and_evaluate(x_raw, y_raw)
    visualize_tsne(x_raw, preds, title="Raw Input t-SNE")

if __name__ == '__main__':
    main()
